import torch
from torch.nn.utils.rnn import pad_sequence


def collate_func(batch_data):
    """
    DataLoader所需的collate_fun函数，将数据处理成tensor形式
    Args:
        batch_data: batch数据

    Returns:

    """
    batch_size = len(batch_data)
    # 如果batch_size为0，则返回一个空字典
    if batch_size == 0:
        return {}
    input_ids_list, token_type_ids_list = [], []
    # 获取一个batch数据中的最大长度
    max_len = max([len(instance["input_ids"]) for instance in batch_data])
    for instance in batch_data:
        # 按照batch中的最大数据长度,对数据进行padding填充
        input_ids_temp = instance["input_ids"]
        input_ids_temp.extend([0] * (max_len - len(instance["input_ids"])))
        token_type_ids_temp = instance["token_type_ids"]
        token_type_ids_temp.extend([0] * (max_len - len(instance["token_type_ids"])))
        # 将list数据转换为tensor数据
        input_ids_list.append(torch.tensor(input_ids_temp, dtype=torch.long))
        token_type_ids_list.append(torch.tensor(token_type_ids_temp, dtype=torch.long))
    return {"input_ids": pad_sequence(input_ids_list, batch_first=True, padding_value=0),
            "token_type_ids": pad_sequence(token_type_ids_list, batch_first=True, padding_value=0)}


def convert_feature(sample, tokenizer=None, max_len=None):
    """
        数据处理函数
        Args:
            sample: 一个字典，包含新闻的正文和新闻的标题，格式为{"content": content, "title": title}
            tokenizer:
            max_len:
        Returns:
        """
    content_id = tokenizer.convert_tokens_to_ids("[Content]")
    title_id = tokenizer.convert_tokens_to_ids("[Title]")
    input_ids = []
    token_type_ids = []
    # 对新闻正文进行tokenizer.tokenize分词
    content_tokens = tokenizer.tokenize(sample["content"])
    # 对新闻标题进行tokenizer.tokenize分词，注意tokenizer中已经将[Space]作为一个分隔符，不会切割成多个字符
    title_tokens = tokenizer.tokenize(sample["title"].replace(" ", "[Space]"))
    # 判断如果正文过长，进行截断
    if len(content_tokens) > max_len - len(title_tokens) - 3:
        content_tokens = content_tokens[:max_len - len(title_tokens) - 3]
    # 生成模型所需的input_ids和token_type_ids
    input_ids.append(tokenizer.cls_token_id)
    token_type_ids.append(content_id)
    input_ids.extend(tokenizer.convert_tokens_to_ids(content_tokens))
    token_type_ids.extend([content_id] * len(content_tokens))
    input_ids.append(tokenizer.sep_token_id)
    token_type_ids.append(content_id)
    input_ids.extend(tokenizer.convert_tokens_to_ids(title_tokens))
    token_type_ids.extend([title_id] * len(title_tokens))
    input_ids.append(tokenizer.sep_token_id)
    token_type_ids.append(title_id)
    # 判断input_ids与token_type_ids长度是否一致
    assert len(input_ids) == len(token_type_ids)
    # 判断input_ids长度是否小于等于最大长度
    if len(input_ids) > max_len:
        input_ids = input_ids[:max_len]
        token_type_ids = token_type_ids[:max_len]
    assert len(input_ids) <= max_len
    return input_ids, token_type_ids
